import os
import torch
import random
from utils import format_esm
import pandas as pd
import csv
from dataset.mine_hard import mine_negative, random_positive
import numpy as np

def collate_fn(batch):
    B = len(batch)
    lengths = [len(b[4]) for b in batch]
    L_max = max([len(b[4]) for b in batch]) + 2
    X = np.zeros([B, L_max, 1024])
    for i, b in enumerate(batch):
        x = b[4]
        x_pad = np.pad(x, [[0,L_max-len(b[4])], [0,0]], 'constant', constant_values=(np.nan, ))
        X[i,:,:] = x_pad
    mask = np.isnan(X)
    X[mask] = 0.
    return batch[0], batch[1], batch[2], batch[3], torch.from_numpy(X).to(dtype=torch.float32), torch.from_numpy(1-mask).to(dtype=torch.float32), lengths


class Triplet_dataset_with_mine_EC(torch.utils.data.Dataset):

    def __init__(self, id_ec, ec_id, mine_neg, data_dir='/state/partition/wzzheng/clean/data/train_valid_split/split100',
                 training_data='split100_train_split_0'):
        self.id_ec = id_ec
        self.ec_id = ec_id
        self.full_list = []
        self.mine_neg = mine_neg
        self.data_dir = data_dir
        self.id_seq_a  = {}
        self.id_seq  = {}
        self.training_data = training_data

        n_label = ['n1', 'n2', 'n3', 'n4', 'n5', 'n6', 'n7', 'n8', 'n9', 'n10', 'n11', 'n12', 'n13','n14','n15','n16']
        self.n_label_dict = {label: i for i, label in enumerate(n_label, 430)}

        for ec in ec_id.keys():
            if '-' not in ec:
                self.full_list.append(ec)

    def __len__(self):
        return len(self.full_list)

    def __getitem__(self, index):
        anchor_ec = self.full_list[index]
        anchor = random.choice(self.ec_id[anchor_ec])

        first_label, second_label, third_label, fourth_label = torch.zeros(7),  torch.zeros(100),  torch.zeros(100),  torch.zeros(450)
        for label in self.id_ec[anchor]:
            label_number = label.split('.')
            first_label[int(label_number[0])-1] = 1
            second_label[int(label_number[1])-1] = 1
            third_label[int(label_number[2])-1] = 1
            if label_number[3].isdigit():
                fourth_label[int(label_number[3])-1] = 1
            else:
                fourth_label[int(self.n_label_dict.get(label_number[3]))] = 1
        all_label = torch.cat((first_label, second_label, third_label, fourth_label))
        pos = random_positive(anchor, self.id_ec, self.ec_id)
        neg = mine_negative(anchor, self.id_ec, self.ec_id, self.mine_neg)
        a = torch.load(self.data_dir + '/esm_data/' + anchor + '.pt')
        p = torch.load(self.data_dir + '/esm_data/' + pos + '.pt')
        n = torch.load(self.data_dir + '/esm_data/' + neg + '.pt')
        return format_esm(a), format_esm(p), format_esm(n)

